import os
import pickle
import random

from config import tran_file, wav_folder, num_valid


def build_vocab(token):
    global char2idx, idx2char
    if not token in char2idx:
        next_index = len(char2idx)
        char2idx[token] = next_index
        idx2char[next_index] = token


def get_blizzard_2013():
    print('loading samples...')

    with open(tran_file, 'r') as file:
        lines = file.readlines()

    num_lines = len(lines)

    samples = []
    for i in range(0, num_lines, 3):
        filename = lines[i].strip()
        filename = '{}/{}.wav'.format(wav_folder, filename)
        # print(filename)
        assert (os.path.isfile(filename))
        text = lines[i + 1].strip()
        samples.append({'audiopath': filename, 'text': text})

    num_samples = len(samples)
    print('num_samples: ' + str(num_samples))

    return samples


if __name__ == "__main__":
    samples = get_blizzard_2013()

    valid = random.sample(samples, num_valid)
    train = []
    for sample in samples:
        if sample not in valid:
            train.append(sample)

    print('num_train: ' + str(len(train)))
    print('num_valid: ' + str(len(valid)))

    with open('train.pkl', 'wb') as file:
        pickle.dump(train, file)
    with open('valid.pkl', 'wb') as file:
        pickle.dump(valid, file)
